# -*- coding: utf-8 -*-
from multiprocessing.pool import CLOSE

from torch.fx.experimental.unification.multipledispatch.dispatcher import source
from ultralytics import YOLO
from PIL import Image
from pathlib import Path

def main():
    model = YOLO(model=r'D:\Clion\code\20241124\runs\detect\train2\weights\best.pt')
    # url = [
    #     "https://img95.699pic.com/photo/50770/1588.jpg_wh300.jpg!/fh/300/quality/90",
    #     "https://news-vod.voc.com.cn/9/2022/11/08/8a21e0592be75430a4b77aca41fb8c0171ff318d1667892249071.jpg?pid=2003269",
    #     "https://x0.ifengimg.com/ucms/2021_24/2F9FB7D7AA3BDB702EC49E552BD3CD452261891B_size2009_w1326_h744.png",
    #     # r"D:\Clion\code\20241124\摔倒参考  50种摔倒方式.mp4"
    #        ]

    def get_all_files(folder_path):
        path = Path(folder_path)
        return [file for file in path.rglob("*") if file.is_file()]

    folder = r"D:\Clion\code\20241124\imgs"
    # url = get_all_files(folder)

    
    # url = r"C:\Users\T-T\Downloads\0.avi"


    class_conf_threshold = {1: 0.6, 3: 0.35}
    results = model.predict(

        # source=url,  # 可以是图片或者视频或者文件夹
        source= 0, # 打开默认摄像头
        # conf=0.05,  # 置信度阈值
        stream=True,
        iou=0.60,  # IoU 阈值
        device=None,  # 使用设备，None 表示自动选择，比如'cpu','0'
        classes=(1, 3),  # 指定要检测的类别
        show=True,  # 是否显示推理图像
        save=True,  # 保存推理结果
        show_labels=True,  # 显示检测的标签
        show_conf=True,  # 显示检测置信度
        show_boxes=True  # 显示检测框
    )
    res = set()
    for result in results:
        detections = []
        for box in result.boxes:
            cls_id = int(box.cls)
            conf = box.conf.item()
            conf_thresh = class_conf_threshold.get(cls_id, 1.1)
            if conf >= conf_thresh:
                detections.append(box)
                res.add(cls_id)
        result.boxes = detections  # 替换为过滤后的检测框
        # result.show()
    return res




    # 打开图片
    # paths = results[0].path.split('\\')[-1]
    # path = r"{0}\{1}".format(results[0].save_dir, paths)
    # image = Image.open(path)
    # image.show()
    # return results[0].boxes
    # res_cls = set()
    # for res in results:
    #     temp = res.boxes.cls.cpu().numpy().astype(int).tolist()
    #     for cls in temp:
    #         res_cls.add(cls)
    # return res_cls
    # results = [result.boxes.cls for result in results]
    # return results
    # return len(results[0].boxes.cls) # 是否有 falling 或者 fall?

if __name__ == '__main__':

    main()
    # print(main())



